# -*- coding: utf-8 -*-
"""
Created on 2018 3.26
@author: hugh
"""
import tensorflow as tf
import os
import cv2
from load_image import gen_test_images,read_train_label
import build_net
import config


tf.flags.DEFINE_string("image_path", "./test_dir/bush.jpg", "Test Data for recognition.")
FLAGS = tf.flags.FLAGS

def predict(test_imgs, IMAGE_HEIGHT, IMAGE_WIDTH, num_classes, net_model="resnet_v2_50",
				num_checkpoints=3,checkpoint_path="model/resnet_v2_50"):
	"""预测函数
	Args:
		test_imgs: 测试图片
		IMAGE_HEIGHT: 测试图片高度
		IMAGE_WIDTH: 测试图片宽度
		num_classes: 图片的类别数量
		net_model: 网络名称
		num_checkpoints: 网络模型保存的数量
		checkpoint_path: 网络模型保存的路径
	Returns：
		index：识别出的图片对应的标签
	"""
	X = tf.placeholder(tf.float32, [None, IMAGE_HEIGHT, IMAGE_WIDTH, 3])
	is_training = tf.placeholder(tf.bool, name='is_training')
	k_prob = tf.placeholder(tf.float32) # dropout

	# 定义模型
	Net = build_net.NetGraph()
	net = Net.get_net(net_model, X, num_classes, k_prob, is_training)

	predict = tf.reshape(net, [-1, num_classes])
	max_idx_p = tf.argmax(predict, 1)

	sess = tf.Session()
	init = tf.global_variables_initializer()
	sess.run(init)

	saver_net = tf.train.Saver(tf.global_variables(), max_to_keep=num_checkpoints)
	if os.path.exists(checkpoint_path):
		saver_net.restore(sess, tf.train.latest_checkpoint(checkpoint_path))
	else:
		print("Please trainning data firstly!")
		sess.close()
		exit(-1)

	index = sess.run([max_idx_p], feed_dict={X: test_imgs,	k_prob: 1.0, is_training: False})
	sess.close()
	return index

if __name__ == '__main__':
	# 准备识别数据
	test_imgs, face_dets = gen_test_images(FLAGS.image_path)
	# 读取训练图片的标签和人名字典
	lable_name_dict = read_train_label()
	num_classes = len(lable_name_dict)	
	# 预测识别
	index = predict(test_imgs,
		  config.IMAGE_HEIGHT, config.IMAGE_WIDTH, num_classes,
		  config.net_model, config.num_checkpoints, config.checkpoint_path)
	labels = index[0].tolist()
	# 打印出识别结果
	for label in labels:
		print(label)
		print("Recognition:{}".format(lable_name_dict[label]))
	# 读取测试图片
	img = cv2.imread(FLAGS.image_path)
	for i, d in enumerate(face_dets):
		#给人脸标记矩形框
		img = cv2.rectangle(img,(d.left(), d.top()), (d.right(), d.bottom()),(0, 255, 0), 1)
		#给人脸标记姓名
		cv2.putText(img, lable_name_dict[labels[i]], (d.left(), d.top() - 5), 
									cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
	# 显示标记后的测试图片
	cv2.imshow('Face recognition result', img) 
	# 等待按任意键
	cv2.waitKey()
	cv2.destroyAllWindows()